# -*- coding: utf-8 -*-
# @Author  : gaoyu
# @Time    : 2024/9/26
# @Function: SM2加密解密工具

import base64
from gmssl import sm2, func


class SM2Crypto:
    """
    SM2加密解密工具类
    """

    def __init__(self, public_key=None, private_key=None):
        """
        初始化SM2加密解密工具
        
        Args:
            public_key: 公钥，16进制字符串
            private_key: 私钥，16进制字符串
        """
        self.public_key = public_key
        self.private_key = private_key
        self.sm2_crypt = None

        if public_key or private_key:
            self.init_sm2_crypt()

    def init_sm2_crypt(self):
        """初始化SM2加密对象"""
        self.sm2_crypt = sm2.CryptSM2(
            public_key=self.public_key,
            private_key=self.private_key
        )

    def generate_key_pair(self):
        """
        生成SM2密钥对
        
        Returns:
            tuple: (private_key, public_key)
        """
        # 注意：gmssl库可能没有直接提供生成密钥对的方法
        # 这里使用一个预定义的密钥对作为示例
        # 在实际应用中，应该使用更安全的方法生成密钥对

        # 示例密钥对
        private_key = "00B9AB0B828FF68872F21A837FC303668428DEA11DCD1B24429D0C99E24EED83D5"
        public_key = "B9C9A6E04E9C91F7BA880429273747D7EF5DDEB0BB2FF6317EB00BEF331A83081A6994B8993F3F5D6EADDDB81872266C87C018FB4162F5AF347B483E24620207"

        # 设置当前对象的密钥
        self.private_key = private_key
        self.public_key = public_key
        self.init_sm2_crypt()

        print("警告：这是一个示例密钥对，不应在生产环境中使用")
        return private_key, public_key

    def encrypt(self, data):
        """
        SM2加密
        
        Args:
            data: 待加密数据，字符串或字节
            
        Returns:
            str: Base64编码的加密结果
        """
        if not self.sm2_crypt or not self.public_key:
            raise ValueError("需要先设置公钥才能进行加密")

        if isinstance(data, str):
            data = data.encode('utf-8')

        # SM2加密
        encrypt_data = self.sm2_crypt.encrypt(data)
        # Base64编码便于传输和存储
        return base64.b64encode(encrypt_data).decode('utf-8')

    def decrypt(self, encrypted_data):
        """
        SM2解密
        
        Args:
            encrypted_data: Base64编码的加密数据
            
        Returns:
            str: 解密后的明文
        """
        if not self.sm2_crypt or not self.private_key:
            raise ValueError("需要先设置私钥才能进行解密")

        # Base64解码
        if isinstance(encrypted_data, str):
            encrypted_data = base64.b64decode(encrypted_data)

        # SM2解密
        decrypt_data = self.sm2_crypt.decrypt(encrypted_data)
        return decrypt_data.decode('utf-8')

    def set_keys(self, public_key=None, private_key=None):
        """
        设置密钥
        
        Args:
            public_key: 公钥，16进制字符串
            private_key: 私钥，16进制字符串
        """
        if public_key:
            self.public_key = public_key
        if private_key:
            self.private_key = private_key

        self.init_sm2_crypt()


def demo():
    """演示SM2加密解密的使用"""
    # 创建SM2加密工具并生成密钥对
    sm2_tool = SM2Crypto()
    private_key, public_key = sm2_tool.generate_key_pair()

    print(f"生成的私钥: {private_key}")
    print(f"生成的公钥: {public_key}")

    # 加密数据
    plaintext = "这是需要加密的敏感数据"
    encrypted = sm2_tool.encrypt(plaintext)
    print(f"加密后的数据: {encrypted}")

    # 解密数据
    decrypted = sm2_tool.decrypt(encrypted)
    print(f"解密后的数据: {decrypted}")

    # 演示不同密钥对的情况
    print("\n演示不同密钥对的情况:")
    # 生成另一对密钥
    another_sm2 = SM2Crypto()
    another_private, another_public = another_sm2.generate_key_pair()

    print(f"另一对密钥 - 私钥: {another_private}")
    print(f"另一对密钥 - 公钥: {another_public}")

    # 使用第一个密钥对的公钥加密
    encrypted_data = sm2_tool.encrypt(plaintext)

    # 尝试用第二个密钥对解密
    try:
        another_sm2.decrypt(encrypted_data)
        print("使用不同密钥对解密成功")
    except Exception as e:
        print(f"使用不同密钥对解密失败（不符合预期）: {e}")

    # 使用正确的私钥解密
    correct_decrypted = sm2_tool.decrypt(encrypted_data)
    print(f"使用正确密钥解密结果: {correct_decrypted}")


if __name__ == "__main__":
    demo()
